/**
 * Copyright 2021 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied。
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * @addtogroup MindSpore
 * @{
 * 
 * @brief 提供MindSpore Lite的模型推理相关接口，该模块下的接口是非线程安全的。
 * 
 * @syscap SystemCapability.Ai.MindSpore
 * @since 9
 */

/**
 * @file model.h
 * @kit MindSporeLiteKit
 * @brief 提供了模型相关接口，可以用于模型创建、模型推理等，该接口是非线程安全的。
 * 
 * @include <mindspore/model.h>
 * @library libmindspore_lite_ndk.so
 * @syscap SystemCapability.Ai.MindSpore
 * @since 9
 */

#ifndef MINDSPORE_INCLUDE_C_API_MODEL_C_H
#define MINDSPORE_INCLUDE_C_API_MODEL_C_H

#include "mindspore/tensor.h"
#include "mindspore/context.h"
#include "mindspore/status.h"

#ifdef __cplusplus
extern "C"
{
#endif
/**
 * @brief 指向模型对象的指针。
 *
 * @since 9
 */
typedef void *OH_AI_ModelHandle;

/**
 * @brief 指向训练配置对象的指针。
 *
 * @since 11
 */
typedef void *OH_AI_TrainCfgHandle;

/**
 * @brief 张量数组结构体，用于存储张量数组指针和张量数组长度。
 *
 * @since 9
 */
typedef struct OH_AI_TensorHandleArray
{
  /** 张量数组长度。 */
  size_t handle_num;
  /** 指向张量数组的指针。 */
  OH_AI_TensorHandle *handle_list;
} OH_AI_TensorHandleArray;

/**
 * @brief 张量维度最大值。
 *
 * @since 9
 */
#define OH_AI_MAX_SHAPE_NUM 32

/**
 * @brief 形状维度大小，预留最大维度是32，当前实际支持的最大维度是8。
 *
 * @since 9
 */
typedef struct OH_AI_ShapeInfo
{
  /** 维度数组长度。 */
  size_t shape_num;
  /** 维度数组。 */
  int64_t shape[OH_AI_MAX_SHAPE_NUM];
} OH_AI_ShapeInfo;

/**
 * @brief 回调函数中传入的算子信息。
 * 
 * @since 9
 */
typedef struct OH_AI_CallBackParam
{
  /** 算子名称。 */
  char *node_name;
  /** 算子类型。 */
  char *node_type;
} OH_AI_CallBackParam;

/**
 * @brief 回调函数指针。\n
 *
 * 该函数指针是用于设置{@link OH_AI_ModelPredict}函数参数中的两个回调函数。\n
 * 该指针指向的函数需要包含三个参数，其中inputs和outputs对应了算子的输入和输出张量，kernel_Info表示当前算子的信息。\n
 * 可以通过回调函数监控算子执行的情况，例如统计算子的执行时间，校验算子的正确性等等。
 *
 * @param inputs 模型输入对应的张量数组结构体。
 * @param outputs 模型输出对应的张量数组结构体。
 * @param kernel_Info 当前算子的信息。
 * @return 回调执行是否成功，若成功返回true，失败则返回false。
 * @since 9
 */
typedef bool (*OH_AI_KernelCallBack)(const OH_AI_TensorHandleArray inputs, const OH_AI_TensorHandleArray outputs,
                                      const OH_AI_CallBackParam kernel_Info);

/**
 * @brief 创建一个模型对象。
 *
 * @return 模型对象指针。
 * @since 9
 */
OH_AI_API OH_AI_ModelHandle OH_AI_ModelCreate(void);

/**
 * @brief 释放一个模型对象。
 *
 * @param model 模型对象指针。
 * @since 9
 */
OH_AI_API void OH_AI_ModelDestroy(OH_AI_ModelHandle *model);

/**
 * @brief 从内存缓冲区加载并编译MindSpore Lite模型。\n
 *
 * 注意，同一个{@link OH_AI_ContextHandle}对象仅能传递给{@link OH_AI_ModelBuild}或者{@link OH_AI_ModelBuildFromFile}一次，如果多次调用该函数需要创建多个不同的{@link OH_AI_ContextHandle}。
 * 
 * @param model 模型对象指针。
 * @param model_data  内存中已经加载的模型数据地址。
 * @param data_size 模型数据的长度。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_context 模型运行时的上下文环境，具体见 {@link OH_AI_ContextHandle}。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 9
 */
OH_AI_API OH_AI_Status OH_AI_ModelBuild(OH_AI_ModelHandle model, const void *model_data, size_t data_size,
                                        OH_AI_ModelType model_type, const OH_AI_ContextHandle model_context);

/**
 * @brief 通过模型文件加载并编译MindSpore Lite模型。\n
 *
 * 注意，同一个{@link OH_AI_ContextHandle}对象仅能传递给{@link OH_AI_ModelBuild}或者{@link OH_AI_ModelBuildFromFile}一次，如果多次调用该函数需要创建多个不同的{@link OH_AI_ContextHandle}。
 * 
 * @param model 模型对象指针。
 * @param model_path 模型文件路径。字符串长度限制跟随文件系统。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_context 模型运行时的上下文环境，具体见 {@link OH_AI_ContextHandle}。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 9
 */
OH_AI_API OH_AI_Status OH_AI_ModelBuildFromFile(OH_AI_ModelHandle model, const char *model_path,
                                                OH_AI_ModelType model_type, const OH_AI_ContextHandle model_context);

/**
 * @brief 调整已编译模型的输入形状。
 *
 * @param model 模型对象指针。
 * @param inputs 模型输入对应的张量数组结构体。
 * @param shape_infos 输入形状信息数组，按模型输入顺序排列的由形状信息组成的数组，模型会按顺序依次调整张量形状。
 * @param shape_info_num 形状信息数组的长度。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 9
 */
OH_AI_API OH_AI_Status OH_AI_ModelResize(OH_AI_ModelHandle model, const OH_AI_TensorHandleArray inputs,
                                          OH_AI_ShapeInfo *shape_infos, size_t shape_info_num);

/**
 * @brief 执行模型推理。
 *
 * @param model 模型对象指针。
 * @param inputs 模型输入对应的张量数组结构体。
 * @param outputs 模型输出对应的张量数组结构体的指针。
 * @param before 模型推理前执行的回调函数。
 * @param after 模型推理后执行的回调函数。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 9
 */
OH_AI_API OH_AI_Status OH_AI_ModelPredict(OH_AI_ModelHandle model, const OH_AI_TensorHandleArray inputs,
                                          OH_AI_TensorHandleArray *outputs, const OH_AI_KernelCallBack before,
                                          const OH_AI_KernelCallBack after);

/**
 * @brief 获取模型的输入张量数组结构体。
 *
 * @param model 模型对象指针。
 * @return 模型输入对应的张量数组结构体。
 * @since 9
 */
OH_AI_API OH_AI_TensorHandleArray OH_AI_ModelGetInputs(const OH_AI_ModelHandle model);

/**
 * @brief 获取模型的输出张量数组结构体。
 *
 * @param model 模型对象指针。
 * @return 模型输出对应的张量数组结构体。
 * @since 9
 */
OH_AI_API OH_AI_TensorHandleArray OH_AI_ModelGetOutputs(const OH_AI_ModelHandle model);

/**
 * @brief 通过张量名获取模型的输入张量。
 *
 * @param model 模型对象指针。
 * @param tensor_name 张量名称。字符串长度跟随系统限制。
 * @return tensor_name所对应的输入张量的张量指针，如果输入中没有该张量则返回空。
 * @since 9
 */
OH_AI_API OH_AI_TensorHandle OH_AI_ModelGetInputByTensorName(const OH_AI_ModelHandle model, const char *tensor_name);

/**
 * @brief 通过张量名获取模型的输出张量。
 *
 * @param model 模型对象指针。
 * @param tensor_name 张量名称。字符串长度跟随系统限制。
 * @return tensor_name所对应的输入张量的张量指针，如果输出中没有该张量则返回空。
 * @since 9
 */
OH_AI_API OH_AI_TensorHandle OH_AI_ModelGetOutputByTensorName(const OH_AI_ModelHandle model, const char *tensor_name);

/**
 * @brief 创建训练配置对象指针，仅用于端侧训练。
 * @return 训练配置对象指针。
 * @since 11
 */
OH_AI_API OH_AI_TrainCfgHandle OH_AI_TrainCfgCreate();

/**
 * @brief 销毁训练配置对象指针，仅用于端侧训练。
 *
 * @param train_cfg 训练配置对象指针。
 * @since 11
 */
OH_AI_API void OH_AI_TrainCfgDestroy(OH_AI_TrainCfgHandle *train_cfg);

/**
 * @brief 获取损失函数的名称列表，仅用于端侧训练。
 *
 * @param train_cfg 训练配置对象指针。
 * @param num 损失函数数量。
 * @return 损失函数的名称列表。
 * @since 11
 */
OH_AI_API char **OH_AI_TrainCfgGetLossName(OH_AI_TrainCfgHandle train_cfg, size_t *num);

/**
 * @brief 设置损失函数的名称列表，仅用于端侧训练。
 *
 * @param train_cfg 训练配置对象指针。
 * @param loss_name 损失函数的名称列表。
 * @param num 损失函数数量。
 * @since 11
 */
OH_AI_API void OH_AI_TrainCfgSetLossName(OH_AI_TrainCfgHandle train_cfg, const char **loss_name, size_t num);

/**
 * @brief 获取训练配置的优化等级，仅用于端侧训练。
 *
 * @param train_cfg 训练配置对象指针。
 * @return 优化等级。
 * @since 11
 */
OH_AI_API OH_AI_OptimizationLevel OH_AI_TrainCfgGetOptimizationLevel(OH_AI_TrainCfgHandle train_cfg);

/**
 * @brief 设置训练配置的优化等级，仅用于端侧训练。
 *
 * @param train_cfg 训练配置对象指针。
 * @param level 优化等级。
 * @since 11
 */
OH_AI_API void OH_AI_TrainCfgSetOptimizationLevel(OH_AI_TrainCfgHandle train_cfg, OH_AI_OptimizationLevel level);

/**
 * @brief 从内存缓冲区加载训练模型，并将模型编译至可在Device上运行的状态，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param model_data 指向存储读入模型文件缓冲区的指针。
 * @param data_size 缓冲区大小。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_context 模型运行时的上下文环境，具体见 {@link OH_AI_ContextHandle}。
 * @param train_cfg 训练配置对象指针。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_TrainModelBuild(OH_AI_ModelHandle model, const void *model_data, size_t data_size,
                                             OH_AI_ModelType model_type, const OH_AI_ContextHandle model_context,
                                             const OH_AI_TrainCfgHandle train_cfg);

/**
 * @brief 根据路径读取加载训练模型，并将模型编译至可在Device上运行的状态，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param model_path 模型文件路径。字符串长度限制跟随文件系统。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_context 模型运行时的上下文环境，具体见 {@link OH_AI_ContextHandle}。
 * @param train_cfg 训练配置对象指针。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_TrainModelBuildFromFile(OH_AI_ModelHandle model, const char *model_path,
                                                     OH_AI_ModelType model_type,
                                                     const OH_AI_ContextHandle model_context,
                                                     const OH_AI_TrainCfgHandle train_cfg);

/**
 * @brief 单步训练模型，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param before 模型推理前执行的回调函数。
 * @param after 模型推理后执行的回调函数。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_RunStep(OH_AI_ModelHandle model, const OH_AI_KernelCallBack before,
                                     const OH_AI_KernelCallBack after);

/**
 * @brief 设置训练的学习率，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param learning_rate 学习率。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ModelSetLearningRate(OH_AI_ModelHandle model, float learning_rate);

/**
 * @brief 获取训练的学习率，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @return 学习率，没有设置优化器时为0.0。
 * @since 11
 */
OH_AI_API float OH_AI_ModelGetLearningRate(OH_AI_ModelHandle model);

/**
 * @brief 获取模型的所有权重Tensors，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @return 模型的所有权重Tensor。
 * @since 11
 */
OH_AI_API OH_AI_TensorHandleArray OH_AI_ModelGetWeights(OH_AI_ModelHandle model);

/**
 * @brief 更新模型的权重Tensor内容，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param new_weights 要更新的权重Tensor。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ModelUpdateWeights(OH_AI_ModelHandle model, const OH_AI_TensorHandleArray new_weights);

/**
 * @brief 获取训练模式。
 *
 * @param model 模型对象指针。
 * @return 表示是否是训练模式。true表示是训练模式，false表示不是训练模式。
 * @since 11
 */
OH_AI_API bool OH_AI_ModelGetTrainMode(OH_AI_ModelHandle model);

/**
 * @brief 设置训练模式，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param train 是否为训练模式。true表示是训练模式，false表示不是训练模式。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ModelSetTrainMode(OH_AI_ModelHandle model, bool train);

/**
 * @brief 设置虚拟batch用于训练，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param virtual_batch_multiplier 虚拟batch乘法器，当设置值小于1时，表示禁用虚拟batch。长度跟随系统限制。
 * @param lr 学习率，默认为-1.0f。
 * @param momentum 动量，默认为-1.0f。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ModelSetupVirtualBatch(OH_AI_ModelHandle model, int virtual_batch_multiplier, float lr,
                                                    float momentum);

/**
 * @brief 导出训练模型，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_file 导出的模型文件路径。字符串长度限制跟随文件系统。
 * @param quantization_type 量化类型。
 * @param export_inference_only 是否导出推理模型。true表示导出推理模型，false表示不导出推理模型。
 * @param output_tensor_name 设置导出模型的输出Tensor，默认为空表示全量导出。
 * @param num 输出Tensor的数量。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ExportModel(OH_AI_ModelHandle model, OH_AI_ModelType model_type, const char *model_file,
                                         OH_AI_QuantizationType quantization_type, bool export_inference_only,
                                         char **output_tensor_name, size_t num);

/**
 * @brief 导出训练模型内存缓存，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param model_data 指向导出模型文件缓冲区的指针。
 * @param data_size 缓冲区大小。
 * @param quantization_type 量化类型。
 * @param export_inference_only 是否导出推理模型。true表示导出推理模型，false表示不导出推理模型。
 * @param output_tensor_name 设置导出模型的输出Tensor，默认为空表示全量导出。
 * @param num 输出Tensor的数量。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ExportModelBuffer(OH_AI_ModelHandle model, OH_AI_ModelType model_type, void *model_data,
                                               size_t *data_size, OH_AI_QuantizationType quantization_type,
                                               bool export_inference_only, char **output_tensor_name, size_t num);

/**
 * @brief 导出模型权重，只能用于micro推理，仅用于端侧训练。
 *
 * @param model 模型对象指针。
 * @param model_type 模型文件类型，具体见{@link OH_AI_ModelType}。
 * @param weight_file 导出的权重文件路径。字符串长度限制跟随文件系统。
 * @param is_inference 是否导出推理模型，当前只支持设置为true。
 * @param enable_fp16 浮点权重是否保存为float16格式。true表示保存为float16格式，false表示不保存为float16格式。
 * @param changeable_weights_name shape可变的权重Tensor名称。
 * @param num shape可变的权重Tensor名称的数量。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 11
 */
OH_AI_API OH_AI_Status OH_AI_ExportWeightsCollaborateWithMicro(OH_AI_ModelHandle model, OH_AI_ModelType model_type,
                                                               const char *weight_file, bool is_inference,
                                                               bool enable_fp16, char **changeable_weights_name,
                                                               size_t num);

/**
 * @brief 加载模型配置文件。
 *
 * @param model 模型对象指针。
 * @param config_path 配置文件路径。字符串长度限制跟随文件系统。
 * @return 枚举类型的状态码{@link OH_AI_Status}，若成功返回OH_AI_STATUS_SUCCESS，失败则返回具体错误码。
 * @since 20
 */
OH_AI_API OH_AI_Status OH_AI_ModelLoadConfig(OH_AI_ModelHandle model, const char *config_path);
															   
#ifdef __cplusplus
}
#endif

/** @} */
#endif // MINDSPORE_INCLUDE_C_API_MODEL_C_H
